########################################################################
# PLS 206 — Week 4 live coding, part 2: reading GAMs, the toolkit, nls
# Grey Monroe, Fall 2026
########################################################################

# install.packages(c("ggplot2", "mgcv"))
library(ggplot2)
library(mgcv)

# fitted values with a 95% CI, for any lm(), glm() or gam()
fit_ci <- function(model, newdata) {
  p <- predict(model, newdata = newdata, se.fit = TRUE)
  inv <- family(model)$linkinv
  data.frame(newdata, fit = inv(as.numeric(p$fit)),
             lwr = inv(as.numeric(p$fit - 1.96 * p$se.fit)),
             upr = inv(as.numeric(p$fit + 1.96 * p$se.fit)))
}

seedlings <- read.csv("https://greymonroe.github.io/PLS_206/data/seedling_temp.csv")
std <- subset(seedlings, variety == "standard")
grid <- data.frame(temp = seq(8.3, 42, by = 0.05))

########################################################################
# 1. The summary, line by line
########################################################################

m_gam <- gam(mass_mg ~ s(temp), data = std, method = "REML")
summary(m_gam)

mean(std$mass_mg)   # = the intercept: smooths are centred
sqrt(m_gam$sig2)    # residual SD, mg (Scale est. is its square)

########################################################################
# 2. Plotting a GAM
########################################################################

plot(m_gam, shade = TRUE)   # partial effect: centred at 0, +-2 SE

# for a figure: back on the scale of y
ci_gam <- fit_ci(m_gam, grid)
ggplot(std, aes(temp, mass_mg)) +
  geom_point() +
  geom_ribbon(data = ci_gam, aes(temp, ymin = lwr, ymax = upr), inherit.aes = FALSE, alpha = 0.3) +
  geom_line(data = ci_gam, aes(temp, fit), colour = "blue") +
  theme_classic()

########################################################################
# 3. Is k big enough?
########################################################################

set.seed(206)
gam.check(m_gam)    # k-index 0.95, p well above 0.05: fine

m4 <- gam(mass_mg ~ s(temp, k = 4), data = std, method = "REML")
set.seed(206)
k.check(m4)         # k-index 0.79, p 0.0025, edf 2.99 of k' = 3: too small

# double k and refit
m8  <- gam(mass_mg ~ s(temp, k = 8),  data = std, method = "REML")
m16 <- gam(mass_mg ~ s(temp, k = 16), data = std, method = "REML")
c(k4 = sum(m4$edf) - 1, k8 = sum(m8$edf) - 1, k16 = sum(m16$edf) - 1)

ci_k <- rbind(cbind(fit_ci(m4, grid), k = "4"),
              cbind(fit_ci(m8, grid), k = "8"),
              cbind(fit_ci(m16, grid), k = "16"))
ggplot(std, aes(temp, mass_mg)) +
  geom_point(colour = "grey60") +
  geom_ribbon(data = ci_k, aes(temp, ymin = lwr, ymax = upr, fill = k), inherit.aes = FALSE, alpha = 0.15) +
  geom_line(data = ci_k, aes(temp, fit, colour = k)) +
  theme_classic()   # 8 and 16 lie on top of each other

########################################################################
# 4. Yes/no: a binomial GAM; one curve per variety
########################################################################

m_dglm <- glm(damaged ~ temp, family = binomial, data = std)
summary(m_dglm)     # slope p = 0.90: a straight line can't make a U

m_dmg <- gam(damaged ~ s(temp), family = binomial, data = std, method = "REML")
summary(m_dmg)

ci_dmg <- fit_ci(m_dmg, grid)   # CI built on the log-odds scale, then plogis()
ggplot(std, aes(temp, damaged)) +
  geom_jitter(width = 0, height = 0.03) +
  geom_ribbon(data = ci_dmg, aes(temp, ymin = lwr, ymax = upr), inherit.aes = FALSE, alpha = 0.3) +
  geom_line(data = ci_dmg, aes(temp, fit), colour = "red") +
  theme_classic()

# both varieties: by = needs the factor in the model too
seedlings$variety <- factor(seedlings$variety)
m_by <- gam(mass_mg ~ variety + s(temp, by = variety), data = seedlings, method = "REML")
summary(m_by)

gv <- expand.grid(temp = seq(8.3, 42, by = 0.1), variety = levels(seedlings$variety))
ci_by <- fit_ci(m_by, gv)
ggplot(seedlings, aes(temp, mass_mg, colour = variety)) +
  geom_point(alpha = 0.5) +
  geom_ribbon(data = ci_by, aes(temp, ymin = lwr, ymax = upr, fill = variety), inherit.aes = FALSE, alpha = 0.3) +
  geom_line(data = ci_by, aes(temp, fit, colour = variety)) +
  theme_classic()

fit_ci(m_by, data.frame(temp = 38, variety = levels(seedlings$variety)))   # 38 degrees

########################################################################
# 5. The field: model the gradient
########################################################################

field <- read.csv("https://greymonroe.github.io/PLS_206/data/field_gradient.csv")
field$genotype <- factor(field$genotype, levels = c("wildtype", "edited"))

ggplot(field, aes(x_m, y_m, fill = yield)) +
  geom_tile() +
  coord_fixed() +
  theme_classic()

m_naive <- lm(yield ~ genotype, data = field)
summary(m_naive)    # p = 0.46

m_sp <- gam(yield ~ genotype + s(x_m, y_m), data = field, method = "REML")
summary(m_sp)       # p = 0.00013

########################################################################
# 6. Maize: a plateau model with nls()
########################################################################

maize <- read.csv("https://greymonroe.github.io/PLS_206/data/maize_nitrogen.csv")

ggplot(maize, aes(N_rate, yield)) +
  geom_point() +
  theme_classic()

# the quadratic's "optimum"
m_mq <- lm(yield ~ N_rate + I(N_rate^2), data = maize)
bm <- coef(m_mq)
-bm[2] / (2 * bm[3])   # 220 kg N/ha

# linear-plateau: starting values read off the plot
m_lp <- nls(yield ~ a + b * pmin(N_rate, Nc), data = maize,
            start = list(a = 5, b = 0.04, Nc = 150))
summary(m_lp)       # Nc = 141 kg N/ha
confint(m_lp)       # 95% CI for Nc: 130 to 154

# a bad start
try(nls(yield ~ a + b * pmin(N_rate, Nc), data = maize,
        start = list(a = 5, b = 0.04, Nc = 400)))   # singular gradient

# both fits on the data (nls band: delta method)
nls_ci <- function(model, newdata) {
  rhs <- formula(model)[[3]]
  f <- function(b) as.numeric(eval(rhs, c(as.list(newdata), as.list(b))))
  b <- coef(model); fit <- f(b)
  J <- sapply(seq_along(b), function(j) {
    h <- 1e-6 * max(1, abs(b[j])); bj <- b; bj[j] <- bj[j] + h
    (f(bj) - fit) / h
  })
  se <- sqrt(rowSums((J %*% vcov(model)) * J))
  q <- qt(0.975, df.residual(model))
  data.frame(newdata, fit = fit, lwr = fit - q * se, upr = fit + q * se)
}

nrates <- data.frame(N_rate = 0:280)
ci_lp <- nls_ci(m_lp, nrates)
ci_mq <- fit_ci(m_mq, nrates)
ggplot(maize, aes(N_rate, yield)) +
  geom_point() +
  geom_ribbon(data = ci_mq, aes(N_rate, ymin = lwr, ymax = upr), inherit.aes = FALSE, alpha = 0.15) +
  geom_line(data = ci_mq, aes(N_rate, fit), colour = "grey50", linetype = 2) +
  geom_ribbon(data = ci_lp, aes(N_rate, ymin = lwr, ymax = upr), inherit.aes = FALSE, alpha = 0.3) +
  geom_line(data = ci_lp, aes(N_rate, fit), colour = "red") +
  geom_vline(xintercept = coef(m_lp)["Nc"], colour = "red", linetype = 3) +
  theme_classic()